import utils.tool_function.experiment_solder_splash as solder_splash
import numpy as np
import cv2
import os


def draw_binary_pic_with_mask(data, target_folder, file, area):
    min_val = data.min()
    max_val = data.max()
    data = (data - min_val) * (255 / (max_val - min_val))
    data = data.astype(np.uint8)
    ret, data = cv2.threshold(data, 50, 255, cv2.THRESH_BINARY)
    cv2.imwrite(
        os.path.join(target_folder, file[:-4] + ".png"), data)


def generate_fly_tin_model(data, target_folder, file, npy_meta):
    min_val = data.min()
    max_val = data.max()
    data = (data - min_val) * (255 / (max_val - min_val))
    data = data.astype(np.uint8)
    ret, data = cv2.threshold(data, 50, 255, cv2.THRESH_BINARY)
    nan_index = data == 0
    npy_meta[nan_index] = np.nan
    result = solder_splash.clean_edge(npy_meta)
    np.save(os.path.join(target_folder, file[:-3] + "npy"), result)
    result_nan = np.isnan(result)
    result[result_nan] = 0
    draw_binary_pic_with_mask(result, target_folder, file[:-4] + "_tin" + ".png", "...")
    return 0


def dfs_count_in_binary_map(binary_data, threshold, area):
    search_flag = binary_data.copy()
    useless_mask = np.ones(binary_data.shape)
    area_44 = []
    while np.sum(search_flag) > 0:
        flag_res = np.where(search_flag > 0)
        dfs_x = flag_res[0][0]
        dfs_y = flag_res[1][0]
        mask_flag = search_flag.copy()
        count, search_flag = solder_splash.dfs_count_points(binary_data, search_flag, dfs_x, dfs_y, 0)
        area_44.append((count, dfs_x, dfs_y))
        if count < threshold:
            useless_mask, search_flag = solder_splash.dfs_points_mask(binary_data, mask_flag, dfs_x, dfs_y,
                                                                      useless_mask)
            binary_data = binary_data * useless_mask
    binary_data = solder_splash.dfs_biggest_area_left(area_44, binary_data)
    return binary_data


def calculate_absolute_tin_diff_clean_nan_with_mask(std_npy, data, mask, area, origin=None, if_show=True,
                                                    target_folder=None, filename=None):
    """
    差值二值化，清除nan对应区域,可视化
    Args:
        std_npy:
        origin: 用于png黑色像素除nan，可不使用
        if_show:
    Returns:

    """
    # print("nan num: ", np.sum(origin))
    # solder_splash.show_3d(data)
    nan_ind1 = np.isnan(data)
    nan_ind2 = np.isnan(std_npy)
    data1 = data.copy()
    data2 = std_npy.copy()
    data1[nan_ind1] = 0
    data2[nan_ind2] = 0
    resul = np.absolute(data1 - data2)
    resul[~mask] = 0
    resul = solder_splash.binary_set(resul)
    resul = dfs_count_in_binary_map(resul, 30, area)
    result = np.sum(resul)
    if (not (filename is None or area is None or target_folder is None)) and result > 200:
        # (area == "43" and result > 200) or (area == "44" and result > 200) or (area == "45" and result > 300)):
        draw_binary_pic_with_mask(resul, target_folder, filename, area)
        generate_fly_tin_model(resul, target_folder, filename, data)
    # resul = clean_nan(origin, resul)
    if if_show:
        solder_splash.show_3d(resul)
    return result


def calculator_tin_diff_with_mask(x_points, y_points, z_points, std_npy, npy, start_point, mask, area, pic=None,
                                  target_folder=None, filename=None):
    """
    利用mask过滤
    Args:
        x_points:
        y_points:
        z_points:
        std_npy:
        npy:
        start_point:

    Returns:

    """
    # origin_nan = np.isnan(npy)
    # 计算底面
    bottom_down, bottom_up, bottom_left = solder_splash.get_bottom_value(x_points, y_points, z_points)
    Z_BOTTOM = (bottom_down + bottom_up + bottom_left) / 3  # PCD面的大致底面高度
    x_points, y_points, z_points = solder_splash.data_filter_by_area(x_points, y_points, z_points,
                                                                     x_min=start_point[1] / solder_splash.RADIO,
                                                                     x_max=(start_point[1] + std_npy.shape[
                                                                         1]) / solder_splash.RADIO,
                                                                     y_min=start_point[0] / solder_splash.RADIO,
                                                                     y_max=(start_point[0] + std_npy.shape[
                                                                         0]) / solder_splash.RADIO)
    real_bottom = solder_splash.get_real_bottom(Z_BOTTOM, z_points)

    # 滑窗计算结果
    result, y, x = solder_splash.move_window(std_npy, npy - real_bottom, start_point)
    print("result: ", str(result), " (", str(x), ", ", str(y), ")")

    # 差异二值化，可视化
    # result = visualize_absolute_tin_diff(std_npy,
    #                                      clean_corner(npy[y:y + std_npy.shape[0], x:x + std_npy.shape[1]] - real_bottom,
    #                                                   0, 0))
    result = calculate_absolute_tin_diff_clean_nan_with_mask(std_npy, solder_splash.clean_corner(
        npy[y:y + std_npy.shape[0], x:x + std_npy.shape[1]] - real_bottom, 0, 0), mask, if_show=False,
                                                             target_folder=target_folder, filename=filename, area=area)
    return result, y, x


def generate_fly_tin(input_npy, target_folder):
    print(input_npy)
    # 清理边缘
    data = solder_splash.clean_height(input_npy)
    # 求pcd板底面深度
    xyz_points = solder_splash.parse_npy_(data)  # 得到x，y，z坐标
    x_points, y_points, z_points = xyz_points

    # 读取标准件区域
    std_4_3 = np.load(solder_splash.STD_4_3)
    std_4_4 = np.load(solder_splash.STD_4_4)
    std_4_5 = np.load(solder_splash.STD_4_5)

    # 亮度图引入，意图清理无效nan点
    pic = solder_splash.clean_light(input_npy[:-4] + "_meta.png")
    pic = cv2.resize(pic, (data.shape[1], data.shape[0]))
    pic = cv2.cvtColor(pic, cv2.COLOR_BGR2GRAY)

    mask43 = np.load("./utils/detection_3d/mask_4-3.npy")
    ind = len(os.listdir(target_folder)) / 3
    ind = int(ind)
    ind = str(ind)
    filename = ind + ".png"
    r43, y43, x43 = calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_3, data, (580, 1100), mask43,
                                                  "43", pic, target_folder=target_folder,
                                                  filename=filename)
    # show_3d(std_4_3, "std_4-3")
    ind = len(os.listdir(target_folder)) / 3
    ind = int(ind)
    ind = str(ind)
    filename = ind + ".png"
    mask44 = np.load("./utils/detection_3d/mask_4-4.npy")
    r44, y44, x44 = calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_4, data, (195, 2050), mask44,
                                                  "44", pic, target_folder=target_folder,
                                                  filename=filename)
    # show_3d(std_4_4, "std_4-4")
    ind = len(os.listdir(target_folder)) / 3
    ind = int(ind)
    ind = str(ind)
    filename = ind + ".png"
    mask45 = np.load("./utils/detection_3d/mask_4-5.npy")
    r45, y45, x45 = calculator_tin_diff_with_mask(x_points, y_points, z_points, std_4_5, data, (10, 2030), mask45, "45",
                                                  pic, target_folder=target_folder,
                                                  filename=filename)

